//
//  29.最小k个数.swift
//  数据结构与算法
//
//  Created by ZERO on 2021/5/21.
//

import Foundation
/*
 题目：给定一个数组，找出其中最小的K个数。例如数组元素是4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4。如果K>数组的长度，那么返回一个空的数组
 思路：1、排序O(nlogn)。2、大顶堆。3、快排思想
 */
func offer_29() {
    print(Solution().getLeastNumbers_Solution3([4,5,1,6,2,7,3,8], 4))
}

fileprivate class Solution {
    // 排序思想
    func getLeastNumbers_Solution1(_ input: [Int],  _ k: Int) -> [Int] {
        // write code here
        guard k <= input.count && k > 0 else { return [] }
        return Array(input.sorted().prefix(k))
    }
    
    // 大顶堆思想
    func getLeastNumbers_Solution2(_ input: [Int],  _ k: Int) -> [Int] {
        // write code here
        guard k <= input.count && k > 0 else { return [] }
        // 建立一个大小为k的大顶堆
        let heap = BigHeap<Int>(k)
        for (i, value) in input.enumerated() {
            if i < k {
                heap.inset(value)
            } else {
                if value <= heap.max()! {
                    heap.removeMax()
                    heap.inset(value)
                }
            }
        }
        return heap.datas.sorted()
    }
    
    // 快排思想
    func getLeastNumbers_Solution3(_ input: [Int],  _ k: Int) -> [Int] {
        // write code here
        guard k <= input.count && k > 0 else { return [] }
        // 分区点的获取
        func partition(arr: inout [Int], l: Int, r: Int) -> Int {
            let pivot = arr[r]
            var i = l
            for j in l..<r {
                if arr[j] < pivot {
                    arr.swapAt(i, j)
                    i += 1
                }
            }
            arr.swapAt(i, r)
            return i
        }
        
        var input = input
        var l = 0
        var r = input.count - 1
        while l <= r {
            let p = partition(arr: &input, l: l, r: r)
            if p < k-1 {
                l = p + 1
            } else if p > k-1 {
                r = p - 1
            } else {
                return Array(input[0..<k]).sorted()
            }
        }
        return []
    }
}

fileprivate class BigHeap<T: Comparable> {
    /// 数组，从下标为1开始存储数据
    private var a: [T?]
    /// 堆可以存储的最大数据个数
    private var n: Int
    /// 堆中已有数据个数
    private var count: Int
    /// 数据源
    var datas: [T] {
        Array(a[1...count].map { $0! })
    }
    
    /// 创建堆
    /// - Parameter capacity: 堆的容量
    init(_ capacity: Int) {
        // 数据从下标为1开始存数组中下标为 i 的节点，左子节点就是下标为 i*2 的节点，右子节点就是下标为 i*2+1 的节点，父节点就是下标为 i/2​ 的节点。
        a = [T?].init(repeating: nil, count: capacity + 1)
        n = capacity
        count = 0
    }
    
    /// 插入元素
    func inset(_ data: T) {
        if count >= n {
            // 堆满了
            return
        }
        count += 1
        a[count] = data
        // 自下而上堆化
        var i = count
        while i/2 > 0 && a[i]! > a[i/2]! {
            a.swapAt(i, i/2)
            i /= 2
        }
    }
    
    /// 删除堆顶元素
    @discardableResult
    func removeMax() -> T? {
        guard count > 0 else { return nil }
        // 把堆顶元素与最后一个元素做交换
        a.swapAt(1, count)
        let temp = a[count]
        a[count] = nil
        count -= 1
        // 自上而下堆化
        heapify(at: 1)
        return temp
    }
    
    /// 返回堆顶元素
    func max() -> T? {
        a[1]
    }
    
    private func heapify(at i: Int) {
        var i = i
        while true {
            var maxIndex = i
            if i*2 <= count && a[i]! < a[i*2]! {
                // 左子节点
                maxIndex = i * 2
            } else if i*2+1 <= count && a[i]! < a[i*2+1]! {
                // 右子节点
                maxIndex = i * 2 + 1
            }
            if maxIndex == i {
                break
            }
            a.swapAt(i, maxIndex)
            i = maxIndex
        }
    }
}

